import collections
import pandas as pd


class SpanF1:
    def __init__(self, method='bio', id2label=None):
        self.metrics = collections.defaultdict(lambda: [0, 0, 0])
        self.id2label = id2label
        self.decoder = self.get_decoder(method)

    def update(self, yt, yp, seq_len):
        for t, p, l in zip(yt, yp, seq_len):
            t = set(self.decoder(t, l))
            p = set(self.decoder(p, l))
            tp_pair = p & t
            fp_pair = p - tp_pair
            fn_pair = t - tp_pair
            self.metrics['gross'][0] += len(tp_pair)
            self.metrics['gross'][1] += len(fp_pair)
            self.metrics['gross'][2] += len(fn_pair)
            for i in tp_pair:
                self.metrics[i[0]][0] += 1
            for i in fp_pair:
                self.metrics[i[0]][1] += 1
            for i in fn_pair:
                self.metrics[i[0]][2] += 1

    @staticmethod
    def get_decoder(method):
        if method == 'bio':
            return decode_bio
        if method == 'bios':
            return decode_bios

    def accumulate(self):
        tmp = []
        index = []
        for k, (tp, fp, fn) in self.metrics.items():
            precision = tp / (tp + fp) if tp + fp > 0 else 0
            recall = tp / (tp + fn) if tp + fn > 0 else 0
            f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0
            tmp.append({'precision': round(precision, 4), 'recall': round(recall, 4), 'f1': round(f1, 4)})
            index.append(k if (self.id2label is None or k == 'gross') else self.id2label[k])
            if k == 'gross':
                main_metrics = f1
        print(pd.DataFrame(tmp, index=index))
        return main_metrics

    def reset(self):
        self.metrics = collections.defaultdict(lambda: [0, 0, 0])


def decode_bio(label, seq_len):
    start = False
    ret = []
    for i, v in enumerate(label[:seq_len]):
        l = (v - 1) // 2
        if v % 2 == 1:
            ret.append([l, i, i])
            start = True
        elif start and v != 0 and l == ret[-1][0]:
            ret[-1][-1] = i
        else:
            start = False
    return [tuple(x) for x in ret]


def decode_bios(label):
    start = False
    ret = []
    for i, v in enumerate(label):
        l = (v - 1) // 3
        if v % 3 == 1:
            ret.append([l, i, i])
            start = True
        elif v % 3 == 2 and start and l == ret[-1][0]:
            ret[-1][-1] = i
        elif v % 3 == 0 and v != 0:
            ret.append([l, i, i])
            start = False
        else:
            start = False
    return [tuple(x) for x in ret]


# def ner_span_f1(y_pred, y_true):
#     all_types = {x['label'] for y in y_true for x in y}
#     y_pred = [{(x['entity'], x['label'], *x['spans']) for x in y} for y in y_pred]
#     y_true = [{(x['entity'], x['label'], *x['spans']) for x in y} for y in y_true]
#     all_types = sorted(list(all_types)) + ['gross']
#     ret = {x: [0, 0, 0] for x in all_types}
#     for y_p, y_t in zip(y_pred, y_true):
#         tp_pair = y_p & y_t
#         fp_pair = y_p - tp_pair
#         fn_pair = y_t - tp_pair
#         ret['gross'][0] += len(tp_pair)
#         ret['gross'][1] += len(fp_pair)
#         ret['gross'][2] += len(fn_pair)
#         for i in tp_pair:
#             ret[i[1]][0] += 1
#         for i in fp_pair:
#             ret[i[1]][1] += 1
#         for i in fn_pair:
#             ret[i[1]][2] += 1
#     tmp = []
#     index = []
#     for k, (tp, fp, fn) in ret.items():
#         precision = tp / (tp + fp) if tp + fp > 0 else 0
#         recall = tp / (tp + fn) if tp + fn > 0 else 0
#         f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0
#         tmp.append({'precision': round(precision, 4), 'recall': round(recall, 4), 'f1': round(f1, 4)})
#         index.append(k)
#         if k == 'gross':
#             main_metrics = f1
#     return main_metrics, pd.DataFrame(tmp, index=index)


# class BIODecoder(OutputDecoder):
#     def __call__(self, labels, texts):
#         start = False
#         output = []
#         for text, label in zip(texts, labels):
#             line = {}
#             entity = collections.defaultdict(list)
#             for i, v in enumerate(label):
#                 if v % 2 == 1:
#                     label = self.id2label[(v - 1) // 2]
#                     entity[label].append([i])
#                     last_start = label
#                     start = True
#                 elif start and v != 0:
#                     label = self.id2label[(v - 1) // 2]
#                     if label == last_start:
#                         entity[label][-1].append(i)
#                 else:
#                     start = False
#             for k, v in entity.items():
#                 # if len(v) > 0:
#                 for j in v:
#                     start = j[0]
#                     end = j[-1]
#                     mapping_start = offset[start][0]
#                     mapping_end = offset[end][1]
#                     entity_text = text[mapping_start:mapping_end]
#                     if entity_text not in line:
#                         line[entity_text] = {"entity": entity_text, 'label': k,
#                                              'spans': [(mapping_start, mapping_end - 1)]}
#                     else:
#                         line[entity_text]['spans'].append((mapping_start, mapping_end - 1))
#             output.append(list(line.values()))
#
#         # elif method == 'span':
#         #     for s, e, t, o in zip(self.start_logits, self.end_logits, self.texts, self.offset):
#         #         s = s.cpu().numpy()[1:-1]
#         #         e = e.cpu().numpy()[1:-1]
#         #         o = o[1:-1]
#         #         for i, s_l in enumerate(s):
#         #             if s_l == 0:
#         #                 continue
#         #             for j, e_l in enumerate(e[i:]):
#         #                 if s_l == e_l:
#         #                     mapping_start = o[i][0]
#         #                     mapping_end = o[i + j][1]
#         #                     entity_text = t[mapping_start:mapping_end]
#         #                     output.append({"entity": entity_text,
#         #                                    'label': self.id2label[s_l - 1],
#         #                                    'spans': [(mapping_start, mapping_end - 1)]})
#         #                     break
#
#         return output
#
#
# class BiAffineDecoder(OutputDecoder):
#     def __call__(self, logits, texts, offsets):
#         output = []
#         labels = self.generate_label(logits)
#         for text, offset, label in zip(texts, offsets, labels):
#             line = {}
#             seq_len = len(offset)
#             indices_mat = torch.triu(torch.ones((seq_len, seq_len)), diagonal=0, out=None).to(labels[0].device)
#             indices = torch.where(indices_mat > 0)
#             logits = label[indices[0], indices[1]].argmax(-1)
#
#             label_mapping = {}
#             for i in range(seq_len):
#                 for j in range(i, seq_len):
#                     label_mapping[len(label_mapping)] = (i, j)
#
#             for idx, value in enumerate(logits):
#                 if value > 0:
#                     start, end = label_mapping[idx]
#                     mapping_start = offset[start][0]
#                     mapping_end = offset[end][1]
#                     entity_text = text[mapping_start:mapping_end]
#                     if entity_text in line:
#                         line[entity_text]['spans'].append((mapping_start, mapping_end - 1))
#                     else:
#                         line[entity_text] = {"entity": entity_text, 'label': self.id2label[value - 1],
#                                              'spans': [(mapping_start, mapping_end - 1)]}
#             output.append(list(line.values()))
#         return output
#
#
# class BIOSDecoder(OutputDecoder):
#     def __call__(self, labels, texts, offsets):
#         start = False
#         output = []
#         for text, offset, label in zip(texts, offsets, labels):
#             line = {}
#             entity = collections.defaultdict(list)
#             for i, v in enumerate(label):
#                 if v % 3 == 1:
#                     label = self.id2label[(v - 1) // 3]
#                     entity[label].append([i])
#                     last_start = label
#                     start = True
#                 elif v % 3 == 2:
#                     label = self.id2label[(v - 1) // 3]
#                     entity[label].append([i])
#                     start = False
#                 elif start and v != 0:
#                     label = self.id2label[(v - 1) // 3]
#                     if label == last_start:
#                         entity[label][-1].append(i)
#                 else:
#                     start = False
#             for k, v in entity.items():
#                 if len(v) > 0:
#                     for j in v:
#                         start = j[0]
#                         end = j[-1]
#                         mapping_start = offset[start][0]
#                         mapping_end = offset[end][1]
#                         entity_text = text[mapping_start:mapping_end]
#                         if entity_text not in line:
#                             line[entity_text] = {"entity": entity_text, 'label': k,
#                                                  'spans': [(mapping_start, mapping_end - 1)]}
#                         else:
#                             line[entity_text]['spans'].append((mapping_start, mapping_end - 1))
#             output.append(list(line.values()))
#         return output


def get_ner_decoder(name):
    ret_dict = {
        'bio': BIODecoder,
        'bios': BIOSDecoder,
        'biaffine': BiAffineDecoder
    }
    if 'bios' in name:
        name = 'bios'
    elif 'bio' in name:
        name = 'bio'
    return ret_dict[name]


if __name__ == '__main__':
    x = [0, 0, 0, 1, 2, 3, 4, 6, 4, 5, 6, 1]
    print(decode_bio(x))

    x = [0, 1, 2, 3, 4, 5, 6, 7, 86, 9, 56, 0]
    print(decode_bios(x))
